import numpy as np
from skimage import measure
import cv2
from skimage.filters import frangi
from sklearn.cluster import KMeans

def background_connect_labels(labels):
    labels = labels[30]
    back_labels = labels[510][510]
    connect_labels = set()
    for i in range(labels.shape[0] - 2):
        for j in range(labels.shape[1] - 2, 0, -1):
            if labels[i][j] != back_labels:
                connect_labels.add(labels[i][j])
                break
        for j in range(labels.shape[1] - 2):
            if labels[i][j] != back_labels:
                connect_labels.add(labels[i][j])
                break
    for j in range(labels.shape[1] - 2):
        for i in range(labels.shape[0] - 2, 0, -1):
            if labels[i][j] != back_labels:
                connect_labels.add(labels[i][j])
                break
        for i in range(labels.shape[0] - 2):
            if labels[i][j] != back_labels:
                connect_labels.add(labels[i][j])
                break

    return back_labels, connect_labels


def make_lungmask(img, display=False):
    #   get the mean of Kmeans center, which used as the threshold
    kmeans = KMeans(n_clusters=4).fit(np.reshape(img, [np.prod(img.shape), 1]))
    centers = sorted(kmeans.cluster_centers_.flatten())
    threshold = np.mean(centers)
    thresh_img = np.where(img < threshold, 1.0, 0.0)  # threshold the image

    #   divide the image after threshold to several label
    labels, num = measure.label(thresh_img, return_num=True)
    props = measure.regionprops(labels)

    #   get the background label and other backgorund-connected labels
    back_label, body_labels = background_connect_labels(labels)

    #   split background labels, the result will be shown in the picture below
    back = labels.copy()
    back[back != back_label] = 3

    #   split Labels connectted with background, the result will be shown in the picture below
    connect = labels.copy()
    for body_label in body_labels:
        connect[connect != body_label] = 3

    #   get Remain part's label, the result will be shown in the picture below
    remain = labels.copy()
    for body_label in body_labels:
        remain[remain == body_label] = back_label
    return remain, back_label


# image: the picture for recognize seed point
# image2: the picture for seed points grow
# threshold: the region max difference,default=50
# seed_num: the seed point we get from original picture, also control the number of slice used for seeking seeds
# start: the first slice num of image we used for seeking seeds
def region_grow(image, image2, threshold=50, seed_num=10, start=16):
    in_region = np.zeros(image.shape, dtype=bool)
    height, width, length = image.shape
    for num in range(seed_num):
        node_x = start + num
        node_y = 0
        node_z = 0
        node_value = 0
        #       get only the brighest point for seed
        for i in range(width):
            for j in range(length):
                temp = image[node_x][i][j]
                if temp > node_value:
                    node_value = temp
                    node_y = i
                    node_z = j

        seed = image2[node_x][node_y][node_z]
        #     set seed point
        in_region[node_x][node_y][node_z] = True
        #       store the avarage value for all seed point
        graysum = seed
        pointsum = 1
        #       use stack to store the pending nodes
        pending = []
        pending.append([node_x, node_y, node_z])

        while len(pending) > 0:
            s = 0
            count = 0
            pend_index = pending.pop()
            for u in range(-1, 2, 1):
                for v in range(-1, 2, 1):
                    for w in range(-1, 2, 1):
                        x = pend_index[0] + u
                        y = pend_index[1] + v
                        z = pend_index[2] + w
                        if (not in_region[x][y][z]) and abs(image2[x][y][z] - seed) <= threshold:
                            in_region[x][y][z] = True
                            pending.append([x, y, z])
                            count += 1
                            s = s + image2[x][y][z]
            pointsum = pointsum + count
            graysum = graysum + s
            seed = graysum / pointsum
    return in_region

def get_lung_pipe_mask(lung, more_precise=True, draw_lung_pipe=False):
    # data = sio.loadmat(path)
    # lung = data['tensor']
    # lung = lung.transpose((2, 0, 1))

    # split lung from body tissues
    lung_mask, back_label = make_lungmask(lung, display=False)

    # allocate weight for lung and other parts
    weight = lung_mask.copy()
    weight[weight != back_label] = 2
    weight[weight == back_label] = 0

    # alloyed the remain lung mask with original picture, augment the backgrounds, get lung part images
    lung_parench = lung.copy() * weight
    lung_parench[weight == 0] = 500
    lung_pipe = lung_parench.copy().astype(np.int16)
    lung_pipe = cv2.GaussianBlur(lung_pipe, (5, 5), 0)

    # calculate the augmented lung
    grayImage = frangi(lung_pipe)
    if not more_precise:
        return_image = grayImage
    else:
        lung_grow = region_grow(grayImage, lung_parench, threshold=60, seed_num=1)
        return_image = lung_grow
    # if draw_lung_pipe:
    #     save_path = './lung3d'
    #     save_lung(return_image, save_path=save_path)
    #     draw_3d(save_path)
    #     return return_image
    # else:
    return return_image
